#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/7/23
# @USER    : Shengji He
# @File    : submit.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
import os
import sys
from SegFlow.inference import BodySegment, TotalBodySegment
from tqdm import tqdm

import SimpleITK as sitk
import pydicom

"""
针对CT图像:
```
TotalSegmentator -i ct.nii.gz -o segmentations
```

"""


def main():
    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    file = r'../weights/Dataset299_body_1559subj/nnUNetTrainer__nnUNetPlans__3d_fullres/fold_0/checkpoint_final.pth'

    save_folder = r'***'
    save_name = 'CT'

    filepath = os.path.join(save_folder, save_name + '.nii.gz')
    save_file = os.path.join(save_folder, 'segmentations')

    segmenter = BodySegment(file, verbose=True, allow_tqdm=True)

    segmenter(filepath, save_file)

    del segmenter


# -------------------------------------------------------------------------------
# for multiple files
# -------------------------------------------------------------------------------
def main_multi_nii():
    from SegFlow.utils import read_csv

    file = r'../weights/Dataset299_body_1559subj/nnUNetTrainer__nnUNetPlans__3d_fullres/fold_0/checkpoint_final.pth'
    segmenter = BodySegment(file, verbose=False, allow_tqdm=False)
    segmenter.build_network()

    folder = r'***'
    csv_file = os.path.join(folder, 'reSegment.csv')
    infos = read_csv(csv_file)
    infos = infos[1:]
    sub_folders = ['***']

    pbar = tqdm(infos)
    for info in pbar:
        case_name = info[1]
        ct_dicom_folder = info[-1]
        flag = False
        for sub_folder in sub_folders:
            if sub_folder in ct_dicom_folder:
                flag = True
                break
        if not flag:
            print(f'{case_name} not found!')
            continue
        save_folder: str = os.path.join(folder, sub_folder)
        case_folder = os.path.join(save_folder, case_name)
        if not os.path.exists(case_folder):
            case_name = f'Subject_{int(case_name):04d}'
            case_folder = os.path.join(save_folder, case_name)
            if not os.path.exists(case_folder):
                print(f'FolderNoteFound: {case_name}')
                continue

        ct_file = os.path.join(case_folder, 'CT.nii.gz')
        if not os.path.exists(ct_file):
            print(f'FileNotFound: {ct_file}')
            continue
        pbar.set_postfix_str(f'processing {case_name}')
        save_file = os.path.join(case_folder, 'segmentations')
        if os.path.exists(save_file):
            continue

        try:
            segmenter(ct_file, save_file)
        except Exception as e:
            print(f'Failed: {case_name};     {e}', file=sys.stderr)

    del segmenter


def main_multi_nii_glob():
    import glob
    import shutil
    from SegFlow.postprocessing import keep_largest_island

    file = r'../weights/Dataset299_body_1559subj/nnUNetTrainer__nnUNetPlans__3d_fullres/fold_0/checkpoint_final.pth'
    segmenter = BodySegment(file, verbose=False, allow_tqdm=False)
    segmenter.build_network()

    folder = r''
    glob_research = os.path.join(folder, '**', 'CT.nii.gz')
    files = glob.glob(glob_research, recursive=True)

    pbar = tqdm(files)
    for ct_file in pbar:
        case_folder = os.path.dirname(ct_file)
        pbar.set_postfix_str(f'processing {case_folder}')
        save_file = os.path.join(case_folder, 'segmentations')
        if os.path.exists(save_file):
            continue

        try:
            segmenter(ct_file, save_file)
        except Exception as e:
            print(f'Failed: {case_folder};     {e}', file=sys.stderr)
        else:
            """
            custom added for binary segmentation for body mask
            """
            file = os.path.join(save_file, 'body.nii.gz')
            save_name = os.path.join(case_folder, 'body.nii.gz')
            img_pred = sitk.ReadImage(file)
            invert_mask = keep_largest_island(img_pred, True)
            img_pred = keep_largest_island(invert_mask, True)
            sitk.WriteImage(img_pred, save_name)
            shutil.rmtree(save_file)
    del segmenter


# -------------------------------------------------------------------------------
# dicom folders from csv
# -------------------------------------------------------------------------------
def SearchDCMFile(folder, labels=None, exclude=None, if_recursion=True):
    """
    Each folder which contain dicom files only exists one case
    """
    reader = sitk.ImageSeriesReader()
    output = []
    if isinstance(folder, str):
        all_folders = set([folder])
    elif isinstance(folder, list):
        all_folders = set(folder)
    if if_recursion:
        out_folder = set()
        for i in all_folders:
            for root, folders, files in os.walk(i):
                for f in folders:
                    out_folder.add(os.path.join(root, f))
        all_folders.update(out_folder)
    for path in all_folders:
        # print('Processing folder: {}'.format(path))
        seriesIDs = reader.GetGDCMSeriesIDs(path)
        if len(seriesIDs) == 0:
            # yield [], '', 0
            continue
        elif len(seriesIDs) > 1:
            print('WARNING: multiple seriesIDs in {}, just SKIP'.format(path))
            # yield [], '', 0
            continue
        else:
            for seriesID in seriesIDs:
                # print('    seriesID: {}'.format(seriesID))
                dicom_names = reader.GetGDCMSeriesFileNames(path, seriesID)
                ds = pydicom.dcmread(dicom_names[0])
                series_description = ''

                # cur_modal = getattr(ds, 'Modality', None)
                # if cur_modal not in ['PT', 'PET']:
                #     print('WARNING: current Modality {}, just SKIP: {}'.format(cur_modal, path))
                #     continue

                if hasattr(ds, 'SeriesDescription'):
                    series_description = ds.SeriesDescription

                if hasattr(ds, 'SeriesType') and len(series_description) == 0:
                    series_description = ds.SeriesType[0]
                if labels is not None and series_description not in labels:
                    continue
                if exclude is not None and series_description in exclude:
                    continue

                output.append([series_description, path, seriesID, dicom_names])
    return output


def main_tb_seg_csv():
    import numpy as np
    import nibabel as nib
    from SegFlow.utils import read_csv

    model_folder = r'../weights'
    segmenter = TotalBodySegment(model_folder, verbose=False, allow_tqdm=False)
    segmenter.task_id = [291]
    segmenter.nr_voxels_thr = 512*512*256
    # segmenter.model_name = 'checkpoint_final_lv.pth'
    segmenter.model_name = 'checkpoint_final.pth'
    # segmenter.perform_everything_on_gpu = False

    folder = r'../output'
    # csv_file = os.path.join(folder, 'manual_confirm_list-20251027-2.csv')
    csv_file = os.path.join(folder, 'seg_plan.csv')
    infos = read_csv(csv_file)
    infos = infos[1:]

    # csv_file = os.path.join(folder, 'manual_confirm_list-20251027-3.csv')
    # infos2 = read_csv(csv_file)
    # infos += infos2[1:]

    # sub_folders = ['Clinicdata_2019-2020', 'PUMC_2015testdata', 'PUMCshare']

    ct_dicom_labels = ['AC_Body', 'AC Body', 'CT_Body', 'CT Body']
    ct_dicom_labels = dict(zip(ct_dicom_labels, range(len(ct_dicom_labels))))

    pbar = tqdm(infos[:])
    for info in pbar:
        sub_folder, case_name, ct_dicom_folder = info
        save_folder: str = os.path.join(folder, sub_folder)
        case_folder = os.path.join(save_folder, case_name)
        # case_name = info[1]
        # ct_dicom_folder = info[14]
        # flag = False
        # for sub_folder in sub_folders:
        #     if sub_folder in ct_dicom_folder:
        #         flag = True
        #         break
        # if not flag:
        #     print(f'{case_name} not found!')
        #     continue
        # save_folder: str = os.path.join(folder, sub_folder)
        # case_folder = os.path.join(save_folder, case_name)
        # if sub_folder == 'PUMCshare':
        #     case_name = f'Subject_{int(case_name):04d}'
        #     case_folder = os.path.join(save_folder, case_name)
        #     # if not os.path.exists(case_folder):
        #     #     print(f'FolderNoteFound: {case_name}')
        #     #     continue
        save_file = os.path.join(case_folder, 'segmentations.nii.gz')
        if os.path.exists(save_file):
            continue
        ct_file = os.path.join(case_folder, 'CT.nii.gz')
        if not os.path.exists(ct_file):
            reader = sitk.ImageSeriesReader()
            seriesIDs = reader.GetGDCMSeriesIDs(ct_dicom_folder)
            if len(seriesIDs) == 0:
                temp_dcm_info = SearchDCMFile(ct_dicom_folder, labels=ct_dicom_labels)
                temp_dcm_info = sorted(temp_dcm_info, key=lambda ele: ct_dicom_labels.get(ele[0], np.inf))
                if len(temp_dcm_info) == 0:
                    msg = f"Failed find ct: {case_name}, skip {ct_dicom_folder}"
                    print(msg)
                    continue
                dicom_names = temp_dcm_info[0][-1]
            elif len(seriesIDs) > 1:
                msg = f"Found multi-ct: {case_name}, skip {ct_dicom_folder}"
                print(msg)
                continue
            else:
                dicom_names = reader.GetGDCMSeriesFileNames(ct_dicom_folder, seriesIDs[0])
            reader.SetFileNames(dicom_names)
            # TODO FIX: wait test, Origin, Spacing, Direction and so on information
            ct_image = reader.Execute()
            ct_array = sitk.GetArrayFromImage(ct_image)

            nii_direction = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)  # FIXME: maybe copy from dicom better?
            pixelSpacing = ct_image.GetSpacing()
            # origin = ct_image.GetOrigin()

            # ct_file = sitk.GetImageFromArray(ct_array)
            # ct_file.SetSpacing(pixelSpacing)
            # ct_file.SetDirection(nii_direction)
            # sitk.WriteImage(ct_file, 'temp.nii.gz')

            origin = (0., 0., 0.)
            # # 3. 构建仿射矩阵
            # # 将方向列表转换为3x3矩阵
            # direction_array = np.array(nii_direction).reshape(3, 3)
            # # 构建仿射矩阵：结合方向、间距和原点
            # affine = np.eye(4)
            # affine[:3, :3] = direction_array @ np.diag(pixelSpacing)
            # affine[:3, 3] = origin

            # 另一种构建仿射矩阵并处理坐标系差异的方法
            affine = np.eye(4)
            affine[:3, :3] = np.array(nii_direction).reshape(3, 3) * pixelSpacing
            affine[:3, 3] = origin
            # 考虑从RAS到LPS的转换：翻转x和y轴
            affine[0, :] = -affine[0, :]  # 翻转x轴
            affine[1, :] = -affine[1, :]  # 翻转y轴
            ct_array_ras = ct_array.transpose(2, 1, 0)  # 从 (z,y,x) 变为 (x,y,z)
            ct_file = nib.Nifti1Image(ct_array_ras, affine)

        pbar.set_postfix_str(f'processing {case_name}')

        # segmenter(ct_file, save_file)

        try:
            segmenter(ct_file, save_file)
        except Exception as e:
            print(f'Failed: {case_name};     {e}', file=sys.stderr)

    del segmenter
    pass


if __name__ == '__main__':
    main_tb_seg_csv()
    # main_multi_nii_glob()
    # main_multi_nii()
    # main()
    print('done')